# -*- coding: utf-8 -*-
# @Time    : 2024/9/9 19:53
# @Author  : Z.J. Zhang
# @Email   : zijingzhang@mail.ustc.edu.cn
# @File    : laser.py
# @Software: PyCharm
import enum
import typing

import matplotlib

import mot.atom

matplotlib.use('tkagg')
import matplotlib.pyplot as plt
import numpy
import scipy.constants as C
from mot.field.base import RectGrid3D
from myutils.geometrical import Plane
from mot.field.base import FieldGenerator
from myutils.filenametools import get_unique_str


class FieldNames:
    @staticmethod
    def get_I_name(laser_name: str): return "I_%s" % laser_name

    @staticmethod
    def get_costheta_name(laser_name: str): return "costheta_%s" % laser_name

    # @staticmethod
    # def get_eta_name(laser_name: str): return "eta_%s" % laser_name

    def __init__(self, laser_names: typing.List[str], ):
        self.laser_names = laser_names
        self.dict_laser_name_to_I = {
            laser_name:
                self.get_I_name(laser_name) for laser_name in self.laser_names
        }
        self.Is = list(self.dict_laser_name_to_I.values())
        self.sumI = "sumI"
        self.dict_laser_name_to_costheta = {
            laser_name: self.get_costheta_name(laser_name) for laser_name in self.laser_names
        }
        self.costhetas = list(self.dict_laser_name_to_costheta.values())


class BeamPolarization(enum.Enum):
    right_handed_circular_polarization = +1
    linear_polarization = 0
    left_handed_circular_polarization = -1


class LaserInfo:

    def __init__(self,
                 k: numpy.ndarray, polarization: typing.Union[int, BeamPolarization] = BeamPolarization(1),
                 name: str = None,
                 detuning_Hz:float = 0.
                 ):
        """

        :param k:
        :param polarization:
        :param name:
        :param frequency:
        :param detuning_Hz:
        """

        self.k = k
        self.polarization = BeamPolarization(polarization)
        if name is None: name = get_unique_str()
        self.name = name
        self.field_name_I = FieldNames.get_I_name(self.name)
        self.field_name_costheta = FieldNames.get_costheta_name(self.name)
        self.detuning_Hz =detuning_Hz
        # self.field_name_eta = FieldNames.get_eta_name(self.name)
    def set_frequency(self,frequency):
        self.set_vacuum_wavelength(C.c/frequency)
    def set_vacuum_wavelength(self,vacuum_wavelength):
        self.k=self.k/numpy.linalg.norm(self.k) * 2*numpy.pi / vacuum_wavelength


def set_frequency(laser_info:LaserInfo,f):
    laser_info.set_frequency(f)
    return laser_info
class LaserFieldGenerator(FieldGenerator):
    def __init__(self, laser_info: LaserInfo, ):
        self.laser_info = laser_info


class NonDiffractionLaserBeam(LaserFieldGenerator):
    """
    Laser beam with Fresnel number $\pi a^2/\lambda z >> 1$, so the diffraction can be ignored.

    菲涅尔数远大于1的激光，可以不考虑经过孔径的衍射 => 束腰几乎不变
    """

    def __init__(self, laser_info: LaserInfo,
                 # k: numpy.ndarray,
                 source_plane: Plane,
                 get_intensity_in_source_plane: typing.Callable[[float, float, float], float],
                 # s
                 ):
        super(NonDiffractionLaserBeam, self).__init__(laser_info)
        self.source_plane = source_plane
        self.__get_intensity_in_source_plane = numpy.vectorize(get_intensity_in_source_plane,
                                                               signature="(),(),()->()")
        self._vec_get_intensity = numpy.vectorize(self.__get_intensity, signature="(),(),()->()")

    def __get_intensity(self, x, y, z):
        pt = numpy.array([x, y, z])
        pt_in_source_plane = self.source_plane.intersection_with_line(pt, self.laser_info.k)
        return self.__get_intensity_in_source_plane(pt_in_source_plane[0], pt_in_source_plane[1],
                                                    pt_in_source_plane[2], )

    def get_intensity_for_multiple_points(self, X: numpy.ndarray, Y: numpy.ndarray, Z: numpy.ndarray):
        """
        :param X: shape (..., 3)
        :param Y: shape (..., 3)
        :param Z: shape (..., 3)
        :return:
        """
        dims = X.ndim
        pts = numpy.array([X, Y, Z]).transpose((*list(range(1, dims + 1)), 0))

        # old_shape = pts.shape
        # pts = pts.reshape((-1, 3))
        pts_in_source_plane = self.source_plane.intersection_with_parallel_lines(pts, self.laser_info.k)
        # pts_in_source_plane = pts_in_source_plane.reshape(old_shape)
        return self.__get_intensity_in_source_plane(pts_in_source_plane[..., 0], pts_in_source_plane[..., 1],
                                                    pts_in_source_plane[..., 2])

    def cache(self, grid: RectGrid3D):
        grid.set_parameter_data({
            (self.laser_info.field_name_I): self.get_intensity_for_multiple_points(grid.X, grid.Y, grid.Z)
        })
        # return grid


if __name__ == '__main__':
    plt.ion()
    aperture = 4e-3
    laser = NonDiffractionLaserBeam(
        LaserInfo(numpy.array([1, 0, 2]), 1),
        Plane(numpy.array([0, 0, -5e-3]), numpy.array([0, 0, 1]), ),
        lambda x, y, z: 2. if x ** 2 + y ** 2 < aperture ** 2 else 0.,
    )
    pts = numpy.array([[1, 0, 1],
                       [2, 0, 1]])
    laser._vec_get_intensity(*pts.T)
    from mot.field.base import default_grid

    grid = default_grid
    from myutils.test_utils import timer


    @timer
    def func2():
        Idata = laser.get_intensity_for_multiple_points(grid.X, grid.Y, grid.Z)
        Idata[laser.source_plane.normal_vec_of_plane[0] * (grid.X - laser.source_plane.point_in_plane[0])
              + laser.source_plane.normal_vec_of_plane[1] * (grid.Y - laser.source_plane.point_in_plane[1])
              + laser.source_plane.normal_vec_of_plane[2] * (grid.Z - laser.source_plane.point_in_plane[2]) < 0] = 0.
        grid.set_parameter_data({
            FieldNames.get_I_name(laser.laser_info.name): Idata
        })


    func2()

    # @timer
    # def func1():
    #     grid.set_parameter_data({
    #         "I": laser._vec_get_intensity(grid.X, grid.Y, grid.Z) # Very slow
    #     })
    #
    #
    # func1()
    X, Z = numpy.meshgrid(grid.xs, grid.zs, indexing='ij')
    plt.figure()


    # @timer
    # def func3():
    #     for i in range(100):
    #         grid.get(grid.X,grid.Y,grid. Z, ['I'])[:, :, 0]
    # func3()

    @timer
    def get_Idata():
        return grid.get(X, 0, Z, [FieldNames.get_I_name(laser.laser_info.name)])[:, :, 0]


    plt.contourf(X, Z, get_Idata())
    plt.scatter(X, Z, s=0.1)
    plt.gca().set_aspect('equal')
